/*
 * Copyright (C) 2014
 *
 * Brick Yang <printfxxx@163.com>
 *
 * This program is free software. You can redistribute it and/or
 * modify it as you like.
 */

/**
 * @file	mtrace.c
 * @brief	Memory alloc/free trace tool
 */

#ifdef MTRACE

#ifdef __KERNEL__
#include <linux/slab.h>
#include <linux/jhash.h>
#include <linux/atomic.h>
#else	/* __KERNEL__ */
#include <errno.h>
#endif	/* __KERNEL__ */

#include "mtrace.h"

#define MTRACE_TBL_HASH_BITS	14

#define mtrace_table_for_each_entry(e, i, tmp, type)						\
	for (tmp = 0, e = &mtrace_table[(type)][(i) & jhash_mask(MTRACE_TBL_HASH_BITS)];	\
	     (tmp < jhash_size(MTRACE_TBL_HASH_BITS)) || (e = NULL, 0);				\
	     tmp++, e = &mtrace_table[(type)][((i) + tmp) & jhash_mask(MTRACE_TBL_HASH_BITS)])

enum {
#ifdef __KERNEL__
	MTRACE_KMEM = 0,
	MTRACE_VMEM,
	MTRACE_DEVM,
	MTRACE_DMACOH,
	MTRACE_SKB,
	MTRACE_PERCPU,
#else
	MTRACE_MEM = 0,
#endif
	MTRACE_TYPE_MAX
};

const char *mtrace_type[MTRACE_TYPE_MAX] = {
#ifdef __KERNEL__
	"kmem",
	"vmem",
	"devm",
	"dmacoh",
	"skb",
	"percpu",
#else
	"mem",
#endif
};

typedef struct mtrace_entry {
	atomic_long_t key_addr;
	const char *name;
	int line;
} mtrace_entry_t;

static mtrace_entry_t mtrace_table[MTRACE_TYPE_MAX][jhash_size(MTRACE_TBL_HASH_BITS)];

static int mtrace_add(long addr, int type, const char *name, int line)
{
	int rc;
	unsigned int i, tmp;
	mtrace_entry_t *e;

	BUG_ON(type >= MTRACE_TYPE_MAX);

	i = jhash(&addr, sizeof(addr), 0);
	mtrace_table_for_each_entry(e, i, tmp, type) {
		if (atomic_long_cmpxchg_acquire(&e->key_addr, 0, addr)) {
			continue;
		}
		e->name = name;
		e->line = line;
		break;
	}

	if (!e) {
		rc = -ENOMEM;
		goto err;
	}
	return 0;
err:
	return rc;
}

static void mtrace_del(long addr, int type, const char *name, int line)
{
	unsigned int i, tmp;
	mtrace_entry_t *e;

	BUG_ON(type >= MTRACE_TYPE_MAX);

	i = jhash(&addr, sizeof(addr), 0);
	mtrace_table_for_each_entry(e, i, tmp, type) {
		if (atomic_long_read_acquire(&e->key_addr) != addr) {
			continue;
		}
		e->name = NULL;
		e->line = -1;
		atomic_long_set_release(&e->key_addr, 0);
		break;
	}

	if (!e) {
		pr_err("%s: failed to find %lx at %s(), line %d\n",
		       mtrace_type[type], addr, name, line);
	}
}

int mtrace_init(void)
{
	int type;
	unsigned int tmp;
	mtrace_entry_t *e;

	for (type = 0; type < MTRACE_TYPE_MAX; type++) {
		mtrace_table_for_each_entry(e, 0, tmp, type) {
			e->name = NULL;
			e->line = -1;
			atomic_long_set_release(&e->key_addr, 0);
		}
	}

	return 0;
}

void mtrace_finish(void)
{
	int type;
	long addr;
	unsigned int tmp;
	mtrace_entry_t *e;

	for (type = 0; type < MTRACE_TYPE_MAX; type++) {
		mtrace_table_for_each_entry(e, 0, tmp, type) {
			if (!(addr = atomic_long_read_acquire(&e->key_addr))) {
				continue;
			}
			pr_warn("%s leak: %lx at %s(), line %d\n",
				mtrace_type[type], addr, e->name, e->line);
		}
	}
}

#ifdef __KERNEL__

int _mtrace_kmem_add(long addr, const char *name, int line)
{
	int rc;

	if (!ZERO_OR_NULL_PTR(addr) && (rc = mtrace_add(addr, MTRACE_KMEM, name, line))) {
		goto err;
	}

	return 0;
err:
	return rc;
}

void _mtrace_kmem_del(long addr, const char *name, int line)
{
	if (!ZERO_OR_NULL_PTR(addr)) {
		mtrace_del(addr, MTRACE_KMEM, name, line);
	}
}

int _mtrace_vmem_add(long addr, const char *name, int line)
{
	int rc;

	if (!ZERO_OR_NULL_PTR(addr) && (rc = mtrace_add(addr, MTRACE_VMEM, name, line))) {
		goto err;
	}

	return 0;
err:
	return rc;
}

void _mtrace_vmem_del(long addr, const char *name, int line)
{
	if (!ZERO_OR_NULL_PTR(addr)) {
		mtrace_del(addr, MTRACE_VMEM, name, line);
	}
}

int _mtrace_devm_add(long addr, const char *name, int line)
{
	int rc;

	if (!ZERO_OR_NULL_PTR(addr) && (rc = mtrace_add(addr, MTRACE_DEVM, name, line))) {
		goto err;
	}

	return 0;
err:
	return rc;
}

void _mtrace_devm_del(long addr, const char *name, int line)
{
	if (!ZERO_OR_NULL_PTR(addr)) {
		mtrace_del(addr, MTRACE_DEVM, name, line);
	}
}

int _mtrace_dmacoh_add(long addr, const char *name, int line)
{
	int rc;

	if (addr && (rc = mtrace_add(addr, MTRACE_DMACOH, name, line))) {
		goto err;
	}

	return 0;
err:
	return rc;
}

void _mtrace_dmacoh_del(long addr, const char *name, int line)
{
	if (addr) {
		mtrace_del(addr, MTRACE_DMACOH, name, line);
	}
}

int _mtrace_skb_add(long addr, const char *name, int line)
{
	int rc;

	if (addr && (rc = mtrace_add(addr, MTRACE_SKB, name, line))) {
		goto err;
	}

	return 0;
err:
	return rc;
}

void _mtrace_skb_del(long addr, const char *name, int line)
{
	if (addr) {
		mtrace_del(addr, MTRACE_SKB, name, line);
	}
}

int _mtrace_percpu_add(long addr, const char *name, int line)
{
	int rc;

	if (addr && (rc = mtrace_add(addr, MTRACE_PERCPU, name, line))) {
		goto err;
	}

	return 0;
err:
	return rc;
}

void _mtrace_percpu_del(long addr, const char *name, int line)
{
	if (addr) {
		mtrace_del(addr, MTRACE_PERCPU, name, line);
	}
}

#else	/* __KERNEL__ */

int _mtrace_mem_add(long addr, const char *name, int line)
{
	int rc;

	if (addr && (rc = mtrace_add(addr, MTRACE_MEM, name, line))) {
		goto err;
	}

	return 0;
err:
	return rc;
}

void _mtrace_mem_del(long addr, const char *name, int line)
{
	if (addr) {
		mtrace_del(addr, MTRACE_MEM, name, line);
	}
}

#endif	/* __KERNEL__ */

#endif	/* _MTRACE_H_ */
